package burp;

import com.intellij.uiDesigner.core.GridConstraints;
import com.intellij.uiDesigner.core.GridLayoutManager;
import com.intellij.uiDesigner.core.Spacer;

import javax.swing.*;
import java.awt.*;
import java.awt.datatransfer.Clipboard;
import java.awt.datatransfer.StringSelection;
import java.awt.event.*;
import java.util.*;


public class BurpGui {
    private JPanel mainPanel;
    private JPanel classPanel;
    private JPanel optionsPanel;
    private JCheckBox enablePollingCheckBox;
    private JButton detectLibraryVersionsButton;
    private JButton clearHistoryAndResetButton;
    private JPanel PollingOption;
    private JTextArea consoleOutput;
    private JList classesFoundList;
    private JList classesNotFoundList;
    private JFormattedTextField pollingIntervalSeconds;
    private JButton pollNow;
    private JButton clearConsole;
    private JButton refreshDNS;
    private JButton copyWordlistButton;
    private TreeSet<String> classesFound = new TreeSet<>(String.CASE_INSENSITIVE_ORDER);
    private TreeMap<String, String> classesNotFound = new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
    private JListRenderer cellRenderer;

    {
// GUI initializer generated by IntelliJ IDEA GUI Designer
// >>> IMPORTANT!! <<<
// DO NOT EDIT OR ADD ANY CODE HERE!
        $$$setupUI$$$();
    }

    /**
     * Method generated by IntelliJ IDEA GUI Designer
     * >>> IMPORTANT!! <<<
     * DO NOT edit this method OR call it in your code!
     *
     * @noinspection ALL
     */
    private void $$$setupUI$$$() {
        mainPanel = new JPanel();
        mainPanel.setLayout(new GridLayoutManager(2, 1, new Insets(16, 16, 16, 16), -1, -1));
        mainPanel.setFocusable(false);
        classPanel = new JPanel();
        classPanel.setLayout(new GridLayoutManager(2, 3, new Insets(0, 0, 0, 0), 0, 0));
        classPanel.setFocusable(false);
        mainPanel.add(classPanel, new GridConstraints(0, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_WANT_GROW, new Dimension(64, 64), null, null, 0, false));
        final JScrollPane scrollPane1 = new JScrollPane();
        scrollPane1.setEnabled(true);
        scrollPane1.setFocusable(false);
        classPanel.add(scrollPane1, new GridConstraints(1, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_WANT_GROW, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_WANT_GROW, new Dimension(60, 16), null, null, 0, false));
        classesFoundList = new JList();
        scrollPane1.setViewportView(classesFoundList);
        final JScrollPane scrollPane2 = new JScrollPane();
        scrollPane2.setFocusable(false);
        classPanel.add(scrollPane2, new GridConstraints(1, 2, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_WANT_GROW, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_WANT_GROW, new Dimension(60, 16), null, null, 0, false));
        classesNotFoundList = new JList();
        classesNotFoundList.setEnabled(true);
        classesNotFoundList.setFocusCycleRoot(false);
        scrollPane2.setViewportView(classesNotFoundList);
        final JLabel label1 = new JLabel();
        label1.setFocusable(true);
        Font label1Font = this.$$$getFont$$$(null, Font.BOLD, -1, label1.getFont());
        if (label1Font != null) label1.setFont(label1Font);
        label1.setText("Classes Found:");
        classPanel.add(label1, new GridConstraints(0, 0, 1, 1, GridConstraints.ANCHOR_WEST, GridConstraints.FILL_NONE, GridConstraints.SIZEPOLICY_FIXED, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        final JLabel label2 = new JLabel();
        Font label2Font = this.$$$getFont$$$(null, Font.BOLD, -1, label2.getFont());
        if (label2Font != null) label2.setFont(label2Font);
        label2.setText("Classes not Found:");
        classPanel.add(label2, new GridConstraints(0, 2, 1, 1, GridConstraints.ANCHOR_WEST, GridConstraints.FILL_NONE, GridConstraints.SIZEPOLICY_FIXED, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        optionsPanel = new JPanel();
        optionsPanel.setLayout(new GridLayoutManager(15, 2, new Insets(0, 0, 0, 0), -1, -1));
        mainPanel.add(optionsPanel, new GridConstraints(1, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, 1, GridConstraints.SIZEPOLICY_FIXED, new Dimension(64, 256), new Dimension(-1, 256), null, 0, true));
        final JScrollPane scrollPane3 = new JScrollPane();
        scrollPane3.setAutoscrolls(true);
        optionsPanel.add(scrollPane3, new GridConstraints(0, 1, 14, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_WANT_GROW, GridConstraints.SIZEPOLICY_FIXED, new Dimension(32, -1), null, new Dimension(-1, 256), 0, false));
        consoleOutput = new JTextArea();
        consoleOutput.setAutoscrolls(true);
        consoleOutput.setEditable(false);
        consoleOutput.setLineWrap(true);
        consoleOutput.setRows(10);
        consoleOutput.setText("");
        consoleOutput.setWrapStyleWord(true);
        scrollPane3.setViewportView(consoleOutput);
        enablePollingCheckBox = new JCheckBox();
        enablePollingCheckBox.setActionCommand("enableDisablePolling");
        enablePollingCheckBox.setEnabled(true);
        enablePollingCheckBox.setSelected(true);
        enablePollingCheckBox.setText("Enable Polling");
        optionsPanel.add(enablePollingCheckBox, new GridConstraints(1, 0, 1, 1, GridConstraints.ANCHOR_NORTHWEST, GridConstraints.FILL_NONE, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        PollingOption = new JPanel();
        PollingOption.setLayout(new GridLayoutManager(1, 4, new Insets(0, 0, 0, 0), -1, -1));
        optionsPanel.add(PollingOption, new GridConstraints(0, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_FIXED, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        final JLabel label3 = new JLabel();
        label3.setText("Poll every");
        PollingOption.add(label3, new GridConstraints(0, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_NONE, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, null, null, null, 0, false));
        final JLabel label4 = new JLabel();
        label4.setText("Seconds");
        PollingOption.add(label4, new GridConstraints(0, 2, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_NONE, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, null, null, null, 0, false));
        pollingIntervalSeconds = new JFormattedTextField();
        pollingIntervalSeconds.setText("30");
        PollingOption.add(pollingIntervalSeconds, new GridConstraints(0, 1, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_NONE, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, new Dimension(32, -1), null, null, 0, false));
        pollNow = new JButton();
        pollNow.setActionCommand("pollNow");
        pollNow.setText("Poll Now");
        PollingOption.add(pollNow, new GridConstraints(0, 3, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_HORIZONTAL, GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        clearConsole = new JButton();
        clearConsole.setActionCommand("clearConsole");
        clearConsole.setText("Clear Console");
        clearConsole.setToolTipText("Clear console to the right");
        optionsPanel.add(clearConsole, new GridConstraints(3, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        refreshDNS = new JButton();
        refreshDNS.setActionCommand("refreshDNS");
        refreshDNS.setText("Refresh DNS Endpoint");
        refreshDNS.setToolTipText("Keep all data, but refresh interaction endpoint for all future generated payloads");
        optionsPanel.add(refreshDNS, new GridConstraints(4, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        clearHistoryAndResetButton = new JButton();
        clearHistoryAndResetButton.setActionCommand("reset");
        clearHistoryAndResetButton.setText("Reset Extension");
        clearHistoryAndResetButton.setToolTipText("Reset will create a new collaborator interaction endpoint and no longer listen for incoming interactions on the previous endpoint. Class lists and the output will be cleared, as well.");
        optionsPanel.add(clearHistoryAndResetButton, new GridConstraints(5, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        detectLibraryVersionsButton = new JButton();
        detectLibraryVersionsButton.setActionCommand("detectLibraryVersion");
        detectLibraryVersionsButton.setText("Detect Library Versions");
        detectLibraryVersionsButton.setToolTipText("Detect library versions uses the bundled wordlists to try and identify library versions on the remote classpath");
        optionsPanel.add(detectLibraryVersionsButton, new GridConstraints(7, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        copyWordlistButton = new JButton();
        copyWordlistButton.setActionCommand("copyWordlist");
        copyWordlistButton.setText("Copy Detect Library Wordlist");
        copyWordlistButton.setToolTipText("Detect library versions uses the bundled wordlists to try and identify library versions on the remote classpath");
        optionsPanel.add(copyWordlistButton, new GridConstraints(8, 0, 1, 1, GridConstraints.ANCHOR_CENTER, GridConstraints.FILL_BOTH, GridConstraints.SIZEPOLICY_CAN_SHRINK | GridConstraints.SIZEPOLICY_CAN_GROW, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
        final JLabel label5 = new JLabel();
        label5.setText("Built-in Analyzer:");
        optionsPanel.add(label5, new GridConstraints(6, 0, 1, 1, GridConstraints.ANCHOR_WEST, GridConstraints.FILL_NONE, GridConstraints.SIZEPOLICY_FIXED, GridConstraints.SIZEPOLICY_FIXED, null, null, null, 0, false));
    }

    /**
     * @noinspection ALL
     */
    private Font $$$getFont$$$(String fontName, int style, int size, Font currentFont) {
        if (currentFont == null) return null;
        String resultName;
        if (fontName == null) {
            resultName = currentFont.getName();
        } else {
            Font testFont = new Font(fontName, Font.PLAIN, 10);
            if (testFont.canDisplay('a') && testFont.canDisplay('1')) {
                resultName = fontName;
            } else {
                resultName = currentFont.getName();
            }
        }
        return new Font(resultName, style >= 0 ? style : currentFont.getStyle(), size >= 0 ? size : currentFont.getSize());
    }

    /**
     * @noinspection ALL
     */
    public JComponent $$$getRootComponent$$$() {
        return mainPanel;
    }

    private class JListRenderer extends DefaultListCellRenderer {
        private final Color CELL_COLOR_1 = Color.decode("#F2F2F2");
        private final Color CELL_COLOR_2 = Color.decode("#FBFBFB");
        private final Color CELL_COLOR_SELECTED = Color.decode("#FFCD81");

        public Component getListCellRendererComponent(JList list, Object value, int index, boolean isSelected, boolean cellHasFocus) {
            Component c = super.getListCellRendererComponent(list, value, index, isSelected, cellHasFocus);
            setText(value.toString());

            if (index % 2 == 0) {
                c.setBackground(CELL_COLOR_1);
            } else {
                c.setBackground(CELL_COLOR_2);
            }

            if (isSelected) {
                c.setBackground(CELL_COLOR_SELECTED);
            }
            return this;
        }

    }

    public Set<String> getClassesFound() {
        return classesFound;
    }

    public Set<String> getClassesNotFound() {
        return classesNotFound.keySet();
    }

    public void reset() {
        classesFound.clear();
        classesNotFound.clear();
        classesFoundList.setListData(new Object[]{});
        classesNotFoundList.setListData(new Object[]{});
        consoleOutput.setText("");
    }

    private class PollingListener implements KeyListener {
        private ActionListener actionListener;
        private BurpGui burpGui;

        public PollingListener(ActionListener actionListener, BurpGui burpGui) {
            this.actionListener = actionListener;
            this.burpGui = burpGui;
        }

        public void keyPressed(KeyEvent keyEvent) {
        }

        public void keyReleased(KeyEvent keyEvent) {
            actionListener.actionPerformed(new ActionEvent(keyEvent.getSource(), keyEvent.getID(), "KEY_TYPED," + burpGui.getPollingInterval()));
        }

        public void keyTyped(KeyEvent keyEvent) {
        }
    }

    public String getPollingInterval() {
        return pollingIntervalSeconds.getText();
    }

    private class CopyListData extends AbstractAction {
        private JList focusOwner = null;

        public CopyListData(JList focusOwner) {
            super("CopyTest");
            this.focusOwner = focusOwner;
        }

        @Override
        public void actionPerformed(ActionEvent e) {
            Clipboard cb = Toolkit.getDefaultToolkit().getSystemClipboard();
            String s = String.join("\n", focusOwner.getSelectedValuesList());
            cb.setContents(new StringSelection(s), null);
        }
    }

    // Input Classes
    public void setupListeners(ActionListener actionListener) {
        pollNow.addActionListener(actionListener);
        pollingIntervalSeconds.addKeyListener(new PollingListener(actionListener, this));
        clearConsole.addActionListener(actionListener);
        enablePollingCheckBox.addActionListener(actionListener);
        detectLibraryVersionsButton.addActionListener(actionListener);
        clearHistoryAndResetButton.addActionListener(actionListener);
        copyWordlistButton.addActionListener(actionListener);
        refreshDNS.addActionListener(actionListener);
        cellRenderer = new JListRenderer();

        // Setup Ctrl-C support
        classesFoundList.getActionMap().put("CopyTest", new CopyListData(classesFoundList));
        classesFoundList.getInputMap(JComponent.WHEN_IN_FOCUSED_WINDOW).put(KeyStroke.getKeyStroke(KeyEvent.VK_C, InputEvent.CTRL_MASK), "CopyTest");

        classesNotFoundList.getActionMap().put("CopyTest", new CopyListData(classesNotFoundList));
        classesNotFoundList.getInputMap(JComponent.WHEN_IN_FOCUSED_WINDOW).put(KeyStroke.getKeyStroke(KeyEvent.VK_C, InputEvent.CTRL_MASK), "CopyTest");

        classesFoundList.setFocusable(true);
        classesNotFoundList.setFocusable(true);

        classesFoundList.setCellRenderer(cellRenderer);
        classesNotFoundList.setCellRenderer(cellRenderer);
    }

    @SuppressWarnings("unchecked")
    private synchronized void refreshClassLists() {
        SwingUtilities.invokeLater(new Runnable() {

            @Override
            public void run() {
                int[] selFound = classesFoundList.getSelectedIndices();
                int[] selNotFound = classesNotFoundList.getSelectedIndices();

                classesFoundList.setListData(classesFound.toArray());
                classesNotFoundList.setListData(classesNotFound.keySet().toArray());

                classesFoundList.setFocusable(true);
                classesNotFoundList.setFocusable(true);

                classesFoundList.setCellRenderer(cellRenderer);
                classesNotFoundList.setCellRenderer(cellRenderer);

                classesFoundList.setSelectedIndices(selFound);
                classesNotFoundList.setSelectedIndices(selNotFound);
            }

        });

    }

    public int getClassesFoundLength() {
        return classesFound.size();
    }

    public synchronized void addClassFound(String classFound) {
        String cls = classesNotFound.remove(classFound
                .replaceAll("d-0-ll", "\\$")
                .replaceAll("d-4-sh", "_")
        );
        if (cls != null) {
            classesFound.add(cls);
            classesNotFound.remove(classFound);
            refreshClassLists();
        }
    }

    public synchronized void addClassNotFound(String classNotFound) {
        if (!classesFound.contains(classNotFound)) {
            classesNotFound.put(classNotFound, classNotFound);
            refreshClassLists();
        }
    }

    public void clearConsole() {
        consoleOutput.setText("");
    }

    public void consolePrintln(String text) {
        consoleOutput.append(text + "\n");
    }

    public boolean isPollingEnabled() {
        return enablePollingCheckBox.isSelected();
    }
}
